#ifndef CONSISTENT_HASH_HPP
#define CONSISTENT_HASH_HPP

#include <functional>

#include "./LoadBalancePolicyIface.h"

namespace LoadBalance
{
    // 通用模板类，禁止使用T!=std::string类型
    template <typename T>
    class ConsistentHash : public LoadBalancePolicyIface<T>
    {
    public:
        ConsistentHash() = delete;
    };

    // 特化模板类，只允许使用T==std::string类型（因为对外接口操作基于std::string）
    // TODO: 优化对外接口操作不只是基于std::string
    template <>
    class ConsistentHash<std::string> : public LoadBalancePolicyIface<std::string>
    {
    public:
        // virtualNodeNum：每个实际节点的虚拟节点数量，增加虚拟节点可改善负载均衡效果（在哈希环分布更广）
        // hashFunc：可选的自定义哈希函数，默认为 std::hash
        ConsistentHash(size_t virtualNodeNum, std::function<size_t(const std::string &)> hashFunc = std::hash<std::string>())
            : virtualNodeNum_(virtualNodeNum), hashFunc_(hashFunc) {}

        // actualNode：要添加的实际节点名称（如服务器地址）。
        // 每个节点会被复制为若干个虚拟节点。每个虚拟节点通过"actualNode + index"计算出唯一的虚拟节点哈希值
        // 这些虚拟节点哈希值存储在哈希环上，并进行排序以便高效查找
        void addNode(const std::string &actualNode)
        {
            std::lock_guard<std::mutex> lock(this->mtx_);

            for (size_t i = 0; i < this->virtualNodeNum_; ++i)
            {
                size_t virtualNodeHash = this->hashFunc_(actualNode + std::to_string(i));

                this->hashCircle_[virtualNodeHash] = actualNode;
                this->sortedHashes_.push_back(virtualNodeHash);
            }
            std::sort(this->sortedHashes_.begin(), this->sortedHashes_.end()); // 默认从小到大
        }

        // actualNode：要移除的实际节点名称
        void delNode(const std::string &actualNode)
        {
            std::lock_guard<std::mutex> lock(this->mtx_);

            for (size_t i = 0; i < this->virtualNodeNum_; ++i)
            {
                size_t virtualNodeHash = this->hashFunc_(actualNode + std::to_string(i));

                this->hashCircle_.erase(virtualNodeHash);
                auto it = std::find(this->sortedHashes_.begin(), this->sortedHashes_.end(), virtualNodeHash);
                if (it != this->sortedHashes_.end())
                {
                    this->sortedHashes_.erase(it);
                }
            }
        }

        // key：要查找的键（如数据的标识符）
        // 根据键的哈希值在哈希环中查找第一个大于等于该值的节点
        // 如果没有找到（即超出哈希环最大值），则回绕到第一个节点
        std::string getNode(const std::string &key)
        {
            std::lock_guard<std::mutex> lock(this->mtx_);

            if (this->hashCircle_.empty())
            {
                std::cout << "No nodes in consistent hash" << std::endl;
                return "";
            }

            size_t keyHash = this->hashFunc_(key);

            // 在已排序的虚拟节点哈希列表中找到第一个大于键哈希值的位置
            auto it = std::upper_bound(this->sortedHashes_.begin(), this->sortedHashes_.end(), keyHash);
            if (it == this->sortedHashes_.end())
            {
                it = this->sortedHashes_.begin(); // 循环回绕取第一个
            }
            return this->hashCircle_[*it]; // 返回虚拟节点哈希值对应的实际节点名称
        }

        // TODO: 这里没用到接口，可以优化接口更通用
        std::string getNext() override
        {
            return "";
        }

    private:
        // 1个实际节点->多个虚拟节点->多个虚拟节点哈希值->1个实际节点
        // 键->键哈希值->第一个大于虚拟节点哈希值（循环回绕取第一个）->实际节点
        size_t virtualNodeNum_;                               // 每个实际节点的虚拟节点数量
        std::function<size_t(const std::string &)> hashFunc_; // 默认或用户自定义的哈希函数

        std::mutex mtx_;
        std::unordered_map<size_t, std::string> hashCircle_; // 虚拟节点哈希值到实际节点名称的映射
        std::vector<size_t> sortedHashes_;                   // 排序的虚拟节点哈希值列表，用于高效查找
    };
} // LoadBalance

#endif // CONSISTENT_HASH_HPP